import os
import json
import torch
import matplotlib.pyplot as plt
import numpy as np


def read_config(config_path):
    """
    description:
       load config from json file
    params:
        @config_path:config flie path
    ret:
        -
    """
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print('{} not exists'.format(config_path))
        exit(-1)


def path_mkdir(base_config):
    """
    description:
       init path and mkdir
    params:
        @base_config:a path dict
    ret:
        -
    """
    assert type(base_config) == dict
    for key in base_config:
        if not os.path.exists(base_config[key]):
            os.mkdir(base_config[key])


def save_data_results(save_path,dict_save_name,ret):
    '''
    Note:
        save experiment data results
    '''
    if not os.path.exists(save_path):
            os.mkdir(save_path)
    
    try:
        with open(os.path.join(save_path,dict_save_name),'w') as f:
            json.dump(ret,f,indent=-1)
            print("Saved training logs to {}".format(os.path.join(save_path,dict_save_name)))
        f.close()
    except:
        print("save data failed")


def save_model_result(save_path,model_name,model):
    '''
    Note:
        save model result
    '''
    if not os.path.exists(save_path):
            os.mkdir(save_path)

    model_name += '.pth'
    torch.save(model.state_dict(), os.path.join(save_path,model_name))
    print("Saved PyTorch Model State to {}".format(os.path.join(save_path,model_name)))


def handle_fl_ret(hlpr):
    save_name:str = hlpr.save_name + '-' + hlpr.task.params.dst
    file_name:str = save_name + '_data.json'

    ret:dict = {}
    ret['T'] = hlpr.params.T
    ret['E'] = hlpr.params.E
    ret['N'] = hlpr.params.N
    ret['C'] = hlpr.params.C
    ret['BS'] = hlpr.params.BS
    ret['optim'] = hlpr.params.optimizer
    ret['lr'] = hlpr.params.lr
    ret['lr_decay'] = hlpr.params.decay
    ret['momentum'] = hlpr.params.momentum
    ret['model'] = hlpr.params.model
    ret['dst'] = hlpr.params.dst
    ret['iid'] = hlpr.params.IID
    ret['exp_round'] = 1
    ret['exp_name'] = save_name

    if hlpr.params.get_addi_params().get('clip'):
        ret['clip_region'] = hlpr.params.get_addi_params().get('clip').region
        ret['clipping_bound'] = hlpr.params.get_addi_params().get('clip').clipping_bound
        
    if hlpr.params.get_addi_params().get('noise'):
        ret['noise_region'] = hlpr.params.get_addi_params().get('noise').region
        ret['noise_mechanism'] = hlpr.params.get_addi_params().get('noise').noise_mechanism
        ret['noise_mean'] = hlpr.params.get_addi_params().get('noise').noise_mean
        ret['noise_multiplier'] = hlpr.params.get_addi_params().get('noise').noise_multiplier

    if hlpr.params.get_addi_params().get('assist'):
        ret['assist_region'] = hlpr.params.get_addi_params().get('assist').region
        ret['dst_res_weight'] = hlpr.params.get_addi_params().get('assist').dst_res_weight
        ret['angle'] = hlpr.params.get_addi_params().get('assist').angle
        
    ret['test_acc'] = hlpr.task.test_acc
    ret['test_loss'] = hlpr.task.test_loss
    #ret['ob_client'] = hlpr.task.ob_client FIXME
    #ret['train_loss'] = hlpr.task.train_loss

    save_data_results(hlpr.save_path,file_name,ret)
    
    return ret